from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("Glibc MCP Server")

"""
面向 glibc 软件中的 nscd（Name Service Cache Daemon）服务管理
所有 MCP tool 列表：

1. glibc_nscd_service_management
   - 用途：管理 nscd 服务，包括启动、重载、状态检查
   - 参数：action (str): 操作类型，支持 "start", "reload", "status"
   - 返回：统一 JSON 结构

2. glibc_nscd_socket_management  
   - 用途：管理 nscd socket，包括启动、状态检查
   - 参数：action (str): 操作类型，支持 "start", "status"
   - 返回：统一 JSON 结构

3. glibc_nscd_socket_detailed
   - 用途：详细的 nscd socket 管理，包括依赖检查、监听状态、重启测试
   - 参数：action (str): 操作类型，支持 "dependencies", "listening", "restart", "connectivity"
   - 返回：统一 JSON 结构
"""

@mcp.tool()
def glibc_nscd_service_management(action: str) -> dict:
    """
    管理 nscd 服务
    
    参数:
        action: 操作类型，支持 "start", "reload", "status"
            - start: 启动 nscd 服务
            - reload: 重载 nscd 服务（先修改配置为调试模式）
            - status: 检查 nscd 服务状态
    
    返回:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    try:
        if action == "start":
            cmd = "systemctl start nscd.service"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        elif action == "reload":
            # 先修改配置
            sed_cmd = "sed -i 's|ExecStart=/usr/sbin/nscd|ExecStart=/usr/sbin/nscd -d|g' /usr/lib/systemd/system/nscd.service"
            subprocess.run(sed_cmd, shell=True, capture_output=True, text=True)
            
            # 重新加载守护进程配置
            daemon_reload_cmd = "systemctl daemon-reload"
            subprocess.run(daemon_reload_cmd, shell=True, capture_output=True, text=True)
            
            # 重载服务
            cmd = "systemctl reload nscd.service"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        elif action == "status":
            cmd = "systemctl status nscd.service"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        else:
            return {
                "success": False,
                "command": f"Unknown action: {action}",
                "exit_code": 1,
                "stdout": "",
                "stderr": f"Unsupported action: {action}. Supported actions: start, reload, status"
            }
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"nscd_service_management {action}",
            "exit_code": 1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def glibc_nscd_socket_management(action: str) -> dict:
    """
    管理 nscd socket
    
    参数:
        action: 操作类型，支持 "start", "status"
            - start: 启动 nscd socket
            - status: 检查 nscd socket 状态
    
    返回:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    try:
        if action == "start":
            cmd = "systemctl start nscd.socket"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        elif action == "status":
            cmd = "systemctl status nscd.socket"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        else:
            return {
                "success": False,
                "command": f"Unknown action: {action}",
                "exit_code": 1,
                "stdout": "",
                "stderr": f"Unsupported action: {action}. Supported actions: start, status"
            }
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"nscd_socket_management {action}",
            "exit_code": 1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def glibc_nscd_socket_detailed(action: str) -> dict:
    """
    详细的 nscd socket 管理功能
    
    参数:
        action: 操作类型，支持 "dependencies", "listening", "restart", "connectivity"
            - dependencies: 检查 nscd socket 的依赖关系
            - listening: 检查 nscd socket 的监听状态
            - restart: 重启 nscd socket 并验证状态
            - connectivity: 测试 nscd socket 的连接性
    
    返回:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    try:
        if action == "dependencies":
            # 检查依赖关系
            cmd1 = "systemctl list-dependencies nscd.socket"
            result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
            
            cmd2 = "systemctl cat nscd.socket"
            result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
            
            combined_stdout = f"=== Dependencies ===\n{result1.stdout}\n\n=== Service File ===\n{result2.stdout}"
            combined_stderr = f"=== Dependencies Errors ===\n{result1.stderr}\n\n=== Service File Errors ===\n{result2.stderr}"
            
            success = (result1.returncode == 0 or result2.returncode == 0)
            
            return {
                "success": success,
                "command": f"{cmd1} && {cmd2}",
                "exit_code": result1.returncode if result1.returncode != 0 else result2.returncode,
                "stdout": combined_stdout,
                "stderr": combined_stderr
            }
            
        elif action == "listening":
            # 检查活动状态
            cmd1 = "systemctl is-active nscd.socket"
            result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
            
            # 检查监听状态
            cmd2 = "ss -lpn | grep nscd || ls -la /var/run/nscd/socket 2>/dev/null || echo 'No nscd listening or socket file found'"
            result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
            
            combined_stdout = f"=== Active Status ===\n{result1.stdout}\n\n=== Listening Status ===\n{result2.stdout}"
            combined_stderr = f"=== Active Status Errors ===\n{result1.stderr}\n\n=== Listening Status Errors ===\n{result2.stderr}"
            
            success = result1.returncode == 0 and result2.returncode == 0
            
            return {
                "success": success,
                "command": f"{cmd1} && {cmd2}",
                "exit_code": result1.returncode if result1.returncode != 0 else result2.returncode,
                "stdout": combined_stdout,
                "stderr": combined_stderr
            }
            
        elif action == "restart":
            # 重启 socket
            cmd1 = "systemctl restart nscd.socket"
            result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
            
            # 等待并检查状态
            import time
            time.sleep(2)
            
            cmd2 = "systemctl status nscd.socket"
            result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
            
            # 检查 socket 文件
            cmd3 = "test -S /var/run/nscd/socket && echo 'Socket file exists' || echo 'Socket file not found'"
            result3 = subprocess.run(cmd3, shell=True, capture_output=True, text=True)
            
            combined_stdout = f"=== Restart Result ===\n{result1.stdout}\n\n=== Status After Restart ===\n{result2.stdout}\n\n=== Socket File Check ===\n{result3.stdout}"
            combined_stderr = f"=== Restart Errors ===\n{result1.stderr}\n\n=== Status Errors ===\n{result2.stderr}\n\n=== Socket File Errors ===\n{result3.stderr}"
            
            success = result1.returncode == 0 and result2.returncode == 0 and result3.returncode == 0
            
            return {
                "success": success,
                "command": f"{cmd1} && {cmd2} && {cmd3}",
                "exit_code": result1.returncode if result1.returncode != 0 else (result2.returncode if result2.returncode != 0 else result3.returncode),
                "stdout": combined_stdout,
                "stderr": combined_stderr
            }
            
        elif action == "connectivity":
            # 测试基本名称服务功能
            cmd = "getent hosts localhost"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            
            return {
                "success": result.returncode == 0,
                "command": cmd,
                "exit_code": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr
            }
            
        else:
            return {
                "success": False,
                "command": f"Unknown action: {action}",
                "exit_code": 1,
                "stdout": "",
                "stderr": f"Unsupported action: {action}. Supported actions: dependencies, listening, restart, connectivity"
            }
            
    except Exception as e:
        return {
            "success": False,
            "command": f"nscd_socket_detailed {action}",
            "exit_code": 1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()